import pygame
import sys
from pygame.locals import *
from enum import Enum


class Piece(Enum):
    EMPTY = 0
    WHITE = 1
    BLACK = 2


def getPosition(a, b):
    x = 28 + a * 40 - 18
    y = 28 + b * 40 - 18
    return x, y


def getRowAndCol(x, y):
    a = int(round((x-28)/40))
    b = int(round((y-28)/40))
    return a, b


# 判断胜负需要的连续棋子数
def continuousPieceCount(a, b, turn, BOARD, BOARDSIZE):
    global DIRECTIONS
    maxPieceCount = 1
    for directionGroup in DIRECTIONS:
        pieceCount = 1
        aCopy, bCopy = a, b
        direction = directionGroup[0]
        aCopy += direction[0]
        bCopy += direction[1]
        while bCopy >= 0 and bCopy < BOARDSIZE and aCopy >= 0 and aCopy < BOARDSIZE and BOARD[bCopy][aCopy] == turn:
            pieceCount += 1
            aCopy += direction[0]
            bCopy += direction[1]
        aCopy, bCopy = a, b
        direction = directionGroup[1]
        aCopy += direction[0]
        bCopy += direction[1]
        while bCopy >= 0 and bCopy < BOARDSIZE and aCopy >= 0 and aCopy < BOARDSIZE and BOARD[bCopy][aCopy] == turn:
            pieceCount += 1
            aCopy += direction[0]
            bCopy += direction[1]
        if pieceCount > maxPieceCount:
            maxPieceCount = pieceCount
    return maxPieceCount


def onClick(pos, turn, BOARD, BOARDSIZE):
    a, b = getRowAndCol(pos[0], pos[1])
    if b >= 0 and b < BOARDSIZE and a >= 0 and a < BOARDSIZE and BOARD[b][a] == Piece.EMPTY:
        BOARD[b][a] = turn
        return True, a, b
    return False, a, b


# 获取某个位置电脑落子的权重
def getWeight(a, b, turn, BOARD, BOARDSIZE):
    weight = 0
    pieceCount = continuousPieceCount(a, b, turn, BOARD, BOARDSIZE)
    if pieceCount >= 5:
        weight += 1000000
    elif pieceCount == 4:
        weight += 8000
    elif pieceCount == 3:
        weight += 200
    elif pieceCount == 2:
        weight += 18
    else:
        weight += 1
    anotherPlayer = Piece.WHITE if turn == Piece.BLACK else Piece.BLACK
    anotherPlayerPieceCount = continuousPieceCount(a, b, anotherPlayer, BOARD, BOARDSIZE)
    if anotherPlayerPieceCount >= 5:
        weight += 100000
    elif anotherPlayerPieceCount == 4:
        weight += 1200
    elif anotherPlayerPieceCount == 3:
        weight += 50
    elif anotherPlayerPieceCount == 2:
        weight += 8
    else:
        weight += 1
    return weight


def computerClick(turn, BOARD, BOARDSIZE):
    maxWeight = 0
    aMaxWeight, bMaxWeight = (0, 0)
    for b in range(BOARDSIZE):
        for a in range(BOARDSIZE):
            if BOARD[b][a] == Piece.EMPTY:
                weight = getWeight(a, b, turn, BOARD, BOARDSIZE)
                if weight > maxWeight:
                    maxWeight = weight
                    aMaxWeight, bMaxWeight = (a, b)
    if BOARD[bMaxWeight][aMaxWeight] != Piece.EMPTY:
        raise Exception('没有位置了')
    BOARD[bMaxWeight][aMaxWeight] = turn
    return aMaxWeight, bMaxWeight


def initBoard(BOARDSIZE):
    BOARD = []
    for b in range(BOARDSIZE):
        row = []
        for a in range(BOARDSIZE):
            row.append(Piece.EMPTY)
        BOARD.append(row)
    return BOARD


def main():
    pygame.init()

    FPS = 30  # frames per second setting
    fpsClock = pygame.time.Clock()

    # set up the window
    DISPLAYSURF = pygame.display.set_mode((615, 615), 0, 32)
    pygame.display.set_caption('五子棋')

    WHITE = (255, 255, 255)
    GREEN = (0, 255, 0)
    bgImg = pygame.image.load('gobang_bg.jpg')
    whiteImg = pygame.image.load('gobang_white.png')
    blackImg = pygame.image.load('gobang_black.png')

    fontObj = pygame.font.Font('simhei.ttf', 50)
    whiteWinTextSurfaceObj = fontObj.render('白棋获胜，单击重新开始', True, GREEN)
    whiteWinTextRectObj = whiteWinTextSurfaceObj.get_rect(center=(308, 308))
    blackWinTextSurfaceObj = fontObj.render('黑棋获胜，单击重新开始', True, GREEN)
    blackWinTextRectObj = blackWinTextSurfaceObj.get_rect(center=(308, 308))

    BOARDSIZE = 15
    BOARD = initBoard(BOARDSIZE)

    gameComplete = False
    turn = Piece.WHITE

    while True:  # run the game loop
        DISPLAYSURF.fill(WHITE)

        DISPLAYSURF.blit(bgImg, (0, 0))
        for b in range(BOARDSIZE):
            for a in range(BOARDSIZE):
                if BOARD[b][a] != Piece.EMPTY:
                    DISPLAYSURF.blit(
                        whiteImg if BOARD[b][a] == Piece.WHITE else blackImg, getPosition(a, b))
        if gameComplete:
            if turn == Piece.WHITE:
                DISPLAYSURF.blit(whiteWinTextSurfaceObj, whiteWinTextRectObj)
            else:
                DISPLAYSURF.blit(blackWinTextSurfaceObj, blackWinTextRectObj)

        for event in pygame.event.get():
            if event.type == QUIT:
                pygame.quit()
                sys.exit()
            elif event.type == MOUSEBUTTONDOWN:  # 鼠标点击事件
                if gameComplete:
                    BOARD = initBoard(BOARDSIZE)
                    gameComplete = False
                    turn = Piece.WHITE
                else:
                    changed, a, b = onClick(event.pos, turn, BOARD, BOARDSIZE)
                    if changed:
                        pieceCount = continuousPieceCount(
                            a, b, turn, BOARD, BOARDSIZE)
                        if pieceCount >= 5:
                            gameComplete = True
                        else:
                            turn = Piece.BLACK if turn == Piece.WHITE else Piece.WHITE
                            # 电脑落子
                            a, b = computerClick(turn, BOARD, BOARDSIZE)
                            pieceCount = continuousPieceCount(
                                a, b, turn, BOARD, BOARDSIZE)
                            if pieceCount >= 5:
                                gameComplete = True
                            else:
                                turn = Piece.BLACK if turn == Piece.WHITE else Piece.WHITE

        pygame.display.update()
        fpsClock.tick(FPS)


DIRECTIONS = (((1, 0), (-1, 0)), ((0, 1), (0, -1)),
              ((1, 1), (-1, -1)), ((1, -1), (-1, 1)))


if __name__ == '__main__':
    main()
